#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
从SQLite获取表结构并生成MySQL的建表语句
"""

import os
import sys
import sqlite3
import argparse

# SQLite到MySQL数据类型的映射
TYPE_MAP = {
    'INTEGER': 'INT',
    'REAL': 'DOUBLE',
    'TEXT': 'TEXT',
    'BLOB': 'BLOB',
    'VARCHAR': 'VARCHAR',
    'BOOLEAN': 'TINYINT(1)',
    'DATETIME': 'DATETIME',
    'DATE': 'DATE',
    'TIME': 'TIME',
    'TIMESTAMP': 'TIMESTAMP',
    'NUMERIC': 'DECIMAL'
}

def get_all_tables(sqlite_cursor):
    """获取SQLite数据库中的所有表名"""
    sqlite_cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
    return [table[0] for table in sqlite_cursor.fetchall()]

def get_table_structure(sqlite_cursor, table_name):
    """获取SQLite表的详细结构"""
    sqlite_cursor.execute(f"PRAGMA table_info({table_name})")
    columns = sqlite_cursor.fetchall()
    return columns

def get_sqlite_foreign_keys(sqlite_cursor, table_name):
    """获取SQLite表的外键约束"""
    sqlite_cursor.execute(f"PRAGMA foreign_key_list({table_name})")
    return sqlite_cursor.fetchall()

def get_sqlite_indexes(sqlite_cursor, table_name):
    """获取SQLite表的索引"""
    sqlite_cursor.execute(f"PRAGMA index_list({table_name})")
    indexes = sqlite_cursor.fetchall()
    
    result = []
    for idx in indexes:
        sqlite_cursor.execute(f"PRAGMA index_info({idx[1]})")
        index_columns = sqlite_cursor.fetchall()
        result.append((idx, index_columns))
    
    return result

def map_sqlite_type_to_mysql(sqlite_type):
    """映射SQLite数据类型到MySQL数据类型"""
    # 提取基本类型（去掉大小限制等）
    base_type = sqlite_type.split('(')[0].upper()
    
    if base_type in TYPE_MAP:
        if '(' in sqlite_type:
            # 保留大小限制
            return TYPE_MAP[base_type] + sqlite_type[len(base_type):]
        return TYPE_MAP[base_type]
    
    # 默认返回TEXT
    return 'TEXT'

def generate_mysql_create_table(table_name, columns, foreign_keys, indexes):
    """生成MySQL创建表的SQL语句"""
    create_table_sql = f"CREATE TABLE IF NOT EXISTS `{table_name}` (\n"
    
    # 添加列定义
    column_defs = []
    primary_key_cols = []
    
    for col in columns:
        col_id, col_name, col_type, col_notnull, col_default, col_pk = col
        
        # 处理数据类型
        mysql_type = map_sqlite_type_to_mysql(col_type)
        
        # 构建列定义
        column_def = f"    `{col_name}` {mysql_type}"
        
        # 处理NOT NULL
        if col_notnull:
            column_def += " NOT NULL"
        
        # 处理默认值
        if col_default is not None:
            if isinstance(col_default, str) and not col_default.isdigit():
                column_def += f" DEFAULT '{col_default}'"
            else:
                column_def += f" DEFAULT {col_default}"
        
        # 记录主键列
        if col_pk:
            primary_key_cols.append(col_name)
        
        column_defs.append(column_def)
    
    # 添加主键约束
    if primary_key_cols:
        primary_key_def = f"    PRIMARY KEY (`{'`, `'.join(primary_key_cols)}`)"
        column_defs.append(primary_key_def)
    
    # 添加外键约束
    for fk in foreign_keys:
        fk_id, fk_seq, fk_table, fk_from, fk_to, fk_on_update, fk_on_delete, fk_match = fk
        fk_def = f"    FOREIGN KEY (`{fk_from}`) REFERENCES `{fk_table}`(`{fk_to}`)"
        
        # 添加ON DELETE规则
        if fk_on_delete:
            fk_def += f" ON DELETE {fk_on_delete}"
        
        # 添加ON UPDATE规则
        if fk_on_update:
            fk_def += f" ON UPDATE {fk_on_update}"
        
        column_defs.append(fk_def)
    
    # 组合所有列定义
    create_table_sql += ",\n".join(column_defs)
    create_table_sql += "\n) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci;"
    
    # 创建索引的SQL语句
    index_sqls = []
    for idx, idx_columns in indexes:
        idx_name, _, is_unique = idx
        
        # 跳过主键索引，因为已经在表定义中添加了
        if idx_name.startswith('sqlite_autoindex_'):
            continue
        
        columns = [col[2] for col in idx_columns]
        
        if is_unique:
            idx_sql = f"CREATE UNIQUE INDEX `{idx_name}` ON `{table_name}` (`{'`, `'.join(columns)}`);"
        else:
            idx_sql = f"CREATE INDEX `{idx_name}` ON `{table_name}` (`{'`, `'.join(columns)}`);"
        
        index_sqls.append(idx_sql)
    
    return create_table_sql, index_sqls

def generate_mysql_schema(sqlite_db_path, output_file=None, tables=None, exclude_tables=None):
    """从SQLite数据库生成MySQL Schema"""
    try:
        # 连接SQLite数据库
        sqlite_conn = sqlite3.connect(sqlite_db_path)
        sqlite_cursor = sqlite_conn.cursor()
        
        # 获取所有表
        all_tables = get_all_tables(sqlite_cursor)
        
        # 过滤要处理的表
        if tables:
            tables_to_process = [t for t in all_tables if t in tables]
        else:
            tables_to_process = all_tables
        
        # 排除指定的表
        if exclude_tables:
            tables_to_process = [t for t in tables_to_process if t not in exclude_tables]
        
        if not tables_to_process:
            print("没有找到需要处理的表")
            return False
        
        # 生成MySQL Schema
        all_sqls = []
        for table in tables_to_process:
            print(f"处理表 '{table}'...")
            
            # 获取表结构
            columns = get_table_structure(sqlite_cursor, table)
            foreign_keys = get_sqlite_foreign_keys(sqlite_cursor, table)
            indexes = get_sqlite_indexes(sqlite_cursor, table)
            
            # 生成建表语句
            create_table_sql, index_sqls = generate_mysql_create_table(table, columns, foreign_keys, indexes)
            
            all_sqls.append(f"-- Table structure for table `{table}`")
            all_sqls.append(create_table_sql)
            
            if index_sqls:
                all_sqls.append(f"-- Indexes for table `{table}`")
                all_sqls.extend(index_sqls)
            
            all_sqls.append("")  # 空行分隔
        
        # 关闭SQLite连接
        sqlite_cursor.close()
        sqlite_conn.close()
        
        # 合并所有SQL语句
        final_sql = "\n".join(all_sqls)
        
        # 写入文件或输出到控制台
        if output_file:
            with open(output_file, 'w', encoding='utf-8') as f:
                f.write(final_sql)
            print(f"MySQL Schema已生成并保存到 '{output_file}'")
        else:
            print(final_sql)
        
        return True
        
    except Exception as e:
        print(f"生成Schema过程中出错: {e}")
        return False

def main():
    parser = argparse.ArgumentParser(description='从SQLite数据库生成MySQL的Schema')
    parser.add_argument('--sqlite', required=True, help='SQLite数据库文件路径')
    parser.add_argument('--output', help='输出文件路径，不指定则输出到控制台')
    parser.add_argument('--tables', help='要处理的表名，以逗号分隔。不指定则处理所有表')
    parser.add_argument('--exclude', help='要排除的表名，以逗号分隔')
    
    args = parser.parse_args()
    
    # 检查SQLite文件是否存在
    if not os.path.exists(args.sqlite):
        print(f"错误: SQLite数据库文件 '{args.sqlite}' 不存在")
        return 1
    
    # 处理表名列表
    tables = args.tables.split(',') if args.tables else None
    exclude_tables = args.exclude.split(',') if args.exclude else None
    
    # 生成Schema
    success = generate_mysql_schema(args.sqlite, args.output, tables, exclude_tables)
    return 0 if success else 1

if __name__ == "__main__":
    sys.exit(main()) 